#
import numpy as np
import matplotlib.pyplot as plt

# 创建一个合成信号用于演示EMD
t = np.linspace(0, 1, 1000)
signal = 2 * np.sin(20 * np.pi * t**2) + 5 * np.cos(120 * np.pi * t) + np.random.normal(0, 0.5, 1000)

# 修正EMD函数中的错误

def simple_emd(data):
    imfs = []
    residue = data.copy()
    while True:
        # 计算极值点
        maxima = np.r_[True, residue[1:] > residue[:-1]] & np.r_[residue[:-1] > residue[1:], True]
        minima = np.r_[True, residue[1:] < residue[:-1]] & np.r_[residue[:-1] < residue[1:], True]
        extrema = np.where(maxima | minima)[0]

        # 如果极值点太少，则停止
        if len(extrema) < 2:
            break

        # 提取极值点的值
        extrema_values = residue[extrema]

        # 分别提取极大值和极小值
        max_values = extrema_values[maxima[extrema]]
        min_values = extrema_values[minima[extrema]]

        # 拟合包络线
        upper_envelope = np.interp(t, t[extrema][maxima[extrema]], max_values)
        lower_envelope = np.interp(t, t[extrema][minima[extrema]], min_values)
        mean_envelope = (upper_envelope + lower_envelope) / 2

        # 提取IMF
        imf = residue - mean_envelope
        residue = residue - imf

        imfs.append(imf)

    imfs.append(residue)
    return imfs

# 再次应用EMD
imfs = simple_emd(signal)

# 绘制结果
plt.figure(figsize=(12, 8))
plt.subplot(len(imfs) + 1, 1, 1)
plt.plot(t, signal, label='Original Signal')
plt.title('Original Signal')

for i, imf in enumerate(imfs):
    plt.subplot(len(imfs) + 1, 1, i + 2)
    plt.plot(t, imf, label=f'IMF {i+1}')
    plt.title(f'IMF {i+1}')

plt.tight_layout()
plt.show()
